package org.shj.spark.udf

import org.apache.spark.sql.expressions.MutableAggregationBuffer
import org.apache.spark.sql.expressions.UserDefinedAggregateFunction
import org.apache.spark.sql.types._
import org.apache.spark.sql.Row
import org.apache.spark.sql.SparkSession

object MyUDAF extends UserDefinedAggregateFunction {
  // 该方法指定具体输入数据的类型
  def inputSchema: StructType = StructType(StructField("inputColumn", LongType) :: Nil)
  
  // 在进行聚合操作的时候所要处理的数据的结果的类型
  def bufferSchema: StructType = {
    StructType(StructField("sum", LongType) :: StructField("count", LongType) :: Nil)
  }
  
  // 最终返回的结果的类型
  def dataType: DataType = DoubleType
  
  // 相同的输入是否总是得到相同的输出。一般情况下都为true
  def deterministic: Boolean = true
  
  // Initializes the given aggregation buffer. The buffer itself is a `Row` that in addition to
  // standard methods like retrieving a value at an index (e.g., get(), getBoolean()), provides
  // the opportunity to update its values. Note that arrays and maps inside the buffer are still
  // immutable.
  /**
   * 在Aggregate 之前每组数据的初始化结果
   */
  def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0) = 0L
    buffer(1) = 0L
  }
    
  // Updates the given aggregation buffer `buffer` with new input data from `input`
  //在进行聚合的时候，每当有新的值进来，对分组后的聚合如何进行计算
  //本地的聚合操作，相当于MapReduce模型中的combiner
  def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    if (!input.isNullAt(0)) {
      buffer(0) = buffer.getLong(0) + input.getLong(0)
      buffer(1) = buffer.getLong(1) + 1
    }
  }
  // Merges two aggregation buffers and stores the updated buffer values back to `buffer1`
  // 最后在分布式节点进行Local Reduce完成后需要进行全局级别的Merge操作
  def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    buffer1(0) = buffer1.getLong(0) + buffer2.getLong(0)
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }
  
  // Calculates the final result
  // 返回的最后的结果 
  def evaluate(buffer: Row): Double = buffer.getLong(0).toDouble / buffer.getLong(1)
  
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder().appName("SparkSqlUDF").master("local").getOrCreate();
    spark.sparkContext.setLogLevel("WARN")
    
    spark.udf.register("myAverage", MyUDAF)

    val df = spark.read.json("E:/workspace/scala/sparkjava/src/main/resources/employees.json")
    df.createOrReplaceTempView("employees")
    df.show()
    
    val result = spark.sql("SELECT myAverage(salary) as average_salary FROM employees")
    result.show()
    
    spark.stop()
  }
}